{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "S9_XTc_lRxEX"
      },
      "source": [
        "##### Copyright 2020 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "4yQK8WLFR1yI"
      },
      "outputs": [],
      "source": [
        "#@title License\n",
        "\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "pW8P5zUZFjPQ"
      },
      "source": [
        "# Train your own Keyword Spotting Model.\n",
        "\n",
        "Before running any cells please enable GPUs for this notebook to speed it up. \n",
        "\n",
        "* *Edit* → *Notebook Settings*\n",
        "* select *GPU* from the *Hardware Accelerator* drop-down\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ndYMcOtK2kqZ"
      },
      "source": [
        "## Download and install Cloud Text-to-Speech SDK\n",
        "\n",
        "Please remember to click the *RESTART RUNTIME* button before proceeding to later cells.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "qE7f5bYK9l2M"
      },
      "outputs": [],
      "source": [
        "!pip install --upgrade google-cloud-texttospeech\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "nwZHcS2seS46"
      },
      "source": [
        "## \u003cfont color=\"red\"\u003eSTOP !\u003c/font\u003e\n",
        "\n",
        "Before proceeding to the next step, did you click on the *RESTART RUNTIME* button in the output?\n",
        "If not please do so now. After the restart, proceed with the next step."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TiJS6BM3R9TJ"
      },
      "source": [
        "## Authenticate and get Cloud credential\n",
        "\n",
        "Next we will authenticate and get Cloud credential in order to call the Google Cloud Text-to-Speech API to synthesize the speech examples.\n",
        "\n",
        "In order to do so, please follow the 5 steps in the 5 cells labeled *A* through *E* below.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "fdL2p659QCh4"
      },
      "source": [
        "### A.  Create and/or selection a Google Cloud project to use.\n",
        "\n",
        "\n",
        "1.  Go to the [Google Cloud project selector](https://console.cloud.google.com/projectselector2/home/dashboard) page.\n",
        "\n",
        "2.  Select a project to use with this notebook, or create a new project by following the next few sub-items. You will need to create a new project if you have none to select from.\n",
        "    * Click on *CREATE PROJECT* near the upper right corner.\n",
        "    * Edit the project name to something more meaningful if you prefer.\n",
        "    * Complete the project creation process by clicking on the *CREATE* button.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-Jai_XKXRwk3"
      },
      "source": [
        "### B.  Enable billing on the Google Cloud project\n",
        "\n",
        "Please note that with the free quota, you will not be charged unless you explicitly grant permission.\n",
        "\n",
        "1.  First, ensure you have a billing account by clicking [here](https://console.cloud.google.com/billing). If you do not have a billing account, you may add one now:\n",
        "    * Click on the *Add billing account* button.\n",
        "    * Follow the instructions in the 2-page workflow. Please note the  text *No auto-payment after free trial ends* on the right.\n",
        "2.  Next, enable billing on the project you selected (or just created) if billing is not already enabled on the project.\n",
        "    * Go to the project page by clicking [here](https://console.cloud.google.com/billing/projects).\n",
        "    * Click on the 3 dots to right of the selected project and select *Change billing* from the menu.\n",
        "    * In the pop-up dialog, select the billing account you want to use. If there is only one billing account available, the message will inform you so and you can directly click on the *SET ACCOUNT* button.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zLPmXkrcSI1_"
      },
      "source": [
        "### C.  Enable the Google Cloud Text-to-Speech API on the project\n",
        "\n",
        "1.  Click [here](https://console.cloud.google.com/flows/enableapi?apiid=texttospeech.googleapis.com) to go to the page for enabling Google Cloud Text-to-Speech API.\n",
        "2. Select the project from the drop-down menu on top (to the right of *Google Cloud Platform*) for registering the application to use Cloud Text-to-Speech API.\n",
        "3.  Click on *Continue*. When the page returns with the message that *The API is enabled*, you may close the page.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "s5c6uyu5SXMA"
      },
      "source": [
        "### D.  Set up authentication\n",
        "\n",
        "1.  Go to the [Create service account key](https://console.cloud.google.com/apis/credentials/serviceaccountkey) page.\n",
        "2. Select the project from the drop-down menu on top (to the right of *Google Cloud Platform*).\n",
        "3.  Select or create a new service account. Follow these steps if you need to create a new service account:\n",
        "    * From the *Service account* drop-down, select *New service account*.\n",
        "    * Enter a *Service account name*.\n",
        "    * For convenience, from the *Role* drop-down, select *Project → Owner*.\n",
        "    * Complete service account creation by clicking the *Create* button at the bottom left.\n",
        "    * You may save the private key for the json file on your computer at this point and close the pop-up dialog.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VFF_GpjvSwSZ"
      },
      "source": [
        "### E.  Copy the email address for the service account and paste it into the form field in the next cell\n",
        "\n",
        "If when running the next cell you encounter the following error:\n",
        "\n",
        "    ERROR: (gcloud.iam.service-accounts.keys.create) RESOURCE_EXHAUSTED: Maximum number of keys on account reached.\n",
        "\n",
        "Go to the [credentials page](https://console.cloud.google.com/apis/credentials) and click on the service account. This takes you to a page listing all the keys associated with the service account, and you can delete unused ones by clicking on the garbage can icon to the right.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "Ki5BD_Dc2v5H"
      },
      "outputs": [],
      "source": [
        "service_account_name = 'YOUR_SERVICE_KEY_ACCOUNT_HERE@SOMETHING.iam.gserviceaccount.com'  #@param {type: \"string\"}\n",
        "\n",
        "!gcloud auth login\n",
        "!gcloud iam service-accounts keys create /tmp/key.json  --iam-account $service_account_name\n",
        "\n",
        "import os\n",
        "os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = '/tmp/key.json'\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "cO99sMUANcRV"
      },
      "source": [
        "## General imports and helper function\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "L9vwjDDvZKYB"
      },
      "outputs": [],
      "source": [
        "#@title Imports\n",
        "%tensorflow_version 1.x\\n\n",
        "from __future__ import division\n",
        "\n",
        "import collections\n",
        "import IPython\n",
        "import functools\n",
        "import math\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import io\n",
        "import os\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "import random\n",
        "import scipy.io.wavfile\n",
        "import tarfile\n",
        "import time\n",
        "import sys\n",
        "\n",
        "from google.colab import output\n",
        "from google.colab import widgets\n",
        "from base64 import b64decode\n",
        "\n",
        "!pip install ffmpeg-python\n",
        "import ffmpeg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "QQARguSpNbO2"
      },
      "outputs": [],
      "source": [
        "#@title Helper function\n",
        "\n",
        "from IPython.display import HTML, display\n",
        "\n",
        "def progress(value, maximum=100):\n",
        "  return IPython.display.HTML(\"\"\"\n",
        "  \u003cprogress value='{value}' max='{max}' style='width: 80%'\u003e{value}\u003c/progress\u003e\n",
        "    \"\"\".format(value=value, max=maximum))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "UTmJE6e8Pa7T"
      },
      "source": [
        "## Get and load the test data\n",
        "\n",
        "The following cell are responsible for getting the data into the colab and creating the embeddings on top which the model is trained.\n",
        "\n",
        "To train a model on a different source of data, replace the next cell with one that copies in your data and change the file scanning cell to scan it correctly.\n",
        "\n",
        "Finally, ensure that global variable MODEL_LABELS is appropriatly set.\n",
        "\n",
        "File scanning is performed to create 2 lists of wav files:\n",
        " * A training file list containing all possible training files. (All files not in testing_list.txt or validation_list.txt)\n",
        " * An evaluation file list that we will use for testing (validation_list.txt)\n",
        "\n",
        "\n",
        "File lists are actually dictionaries with the following structure:\n",
        "\n",
        "```\n",
        "{'keyword1': ['path/to/word1/example1.wav', path/to/word1/example2.wav'],\n",
        " 'keyword2': ['path/to/word2/example1.wav', path/to/word2/example2.wav'],\n",
        " ...\n",
        " 'negativ': ['path/to/negativ_example1.wav', path/to/negativ_example2.wav']}\n",
        " ```\n",
        "\n",
        "The subsequent cells assume that the file lists are stored in the  variables: *all_eval_example_files* and *all_train_example_files*.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "dkTXhDVdPsgj"
      },
      "outputs": [],
      "source": [
        "#@title Download and extract the speech commands data set\n",
        "data_source = \"http://download.tensorflow.org/data/speech_commands_v0.02.tar.gz\"\n",
        "data_dest_dir = \"speech_commands_v0.02\"\n",
        "test_list = data_dest_dir + \"/testing_list.txt\"\n",
        "valid_list = data_dest_dir + \"/validation_list.txt\"\n",
        "\n",
        "TARGET_WORDS = 'yes,no,up,down,left,right,on,off,stop,go'\n",
        "ALL_WORDS = 'backward,bed,bird,cat,dog,down,eight,five,follow,forward,four,go,' + 'happy,house,learn,left,marvin,nine,no,off,on,one,right,seven,sheila,six,stop,' + 'three,tree,two,up,visual,wow,yes,zero'\n",
        "\n",
        "# Note: This example colab doesn't train the silence output. \n",
        "MODEL_LABELS = ['negative', 'silence'] + TARGET_WORDS.split(',')\n",
        "\n",
        "!wget http://download.tensorflow.org/data/speech_commands_v0.02.tar.gz\n",
        "\n",
        "print(\"extracting tar archive.. this may take a few minutes.\")\n",
        "\n",
        "if not os.path.exists(data_dest_dir):\n",
        "  os.makedirs(data_dest_dir)\n",
        "#tarfile.open(\"speech_commands_v0.02.tar.gz\", 'r:gz').extractall(data_dest_dir)\n",
        "file_count = 0\n",
        "progress_bar = IPython.display.display(progress(0, 100), display_id=True)\n",
        "\n",
        "with tarfile.open(\"speech_commands_v0.02.tar.gz\", 'r:gz') as speech_commands_tar:\n",
        "  for member_info in speech_commands_tar.getmembers():\n",
        "    if file_count % 100 == 0:\n",
        "      progress_bar.update(progress(100 * file_count/105800))\n",
        "    speech_commands_tar.extract(member_info, data_dest_dir)\n",
        "    file_count+=1\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "r4x--vCLP17M"
      },
      "outputs": [],
      "source": [
        "#@title Scan files\n",
        "\n",
        "progress_bar = IPython.display.display(progress(0, 100), display_id=True)\n",
        "print(\"loading filelists from: %s \" % data_dest_dir)\n",
        "\n",
        "def get_test_files(word, test_file_names):\n",
        "  word_dir = os.path.join(data_dest_dir, word)\n",
        "  all_word_files = [\n",
        "      os.path.join(word, f)\n",
        "      for f in os.listdir(word_dir)\n",
        "      if os.path.isfile(os.path.join(word_dir, f))\n",
        "  ]\n",
        "  word_test_files = [f for f in all_word_files if f in test_files]\n",
        "  random.shuffle(word_test_files)\n",
        "  return word_test_files\n",
        "\n",
        "test_files = [line.rstrip() for line in open(test_list, encoding=\"ISO-8859-1\")]\n",
        "\n",
        "all_eval_example_files = collections.defaultdict(list)\n",
        "\n",
        "all_word_list = ALL_WORDS.split(',')\n",
        "\n",
        "word_count = 0\n",
        "for word in all_word_list:\n",
        "  if word in MODEL_LABELS:\n",
        "    label = word\n",
        "  else:\n",
        "    label = \"negative\"\n",
        "  eval_files = get_test_files(word, test_files)\n",
        "  all_eval_example_files[label].extend(eval_files)\n",
        "  if progress is not None:\n",
        "    word_count += 1\n",
        "    progress_bar.update(progress(100 * word_count/len(all_word_list)))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "HLxWDDs32Z11"
      },
      "source": [
        "## Prepare synthetic data\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MHi7YYdCQBux"
      },
      "source": [
        "Here the wav files from both evaluation and training sets are:\n",
        " * Opened and decoded.\n",
        " * Loudness normalized.\n",
        " * Passed through the embedding model to create embeddings.\n",
        " * Added to a data structure that let's us change the balance between negative and labeled outputs.\n",
        "\n",
        "\n",
        " resulting in two objects *eval_data* and *train_data*.\n",
        " "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "BrgMOJspO2b_"
      },
      "outputs": [],
      "source": [
        "#@title Define list of cloud voices of interest\n",
        "\n",
        "standard_voices = [\n",
        "  # Australia\n",
        "  \"en-AU-Standard-A\", \"en-AU-Standard-B\", \"en-AU-Standard-C\", \"en-AU-Standard-D\",\n",
        "  # UK\n",
        "  \"en-GB-Standard-A\", \"en-GB-Standard-B\", \"en-GB-Standard-C\", \"en-GB-Standard-D\",\n",
        "  # India\n",
        "  \"en-IN-Standard-A\", \"en-IN-Standard-B\", \"en-IN-Standard-C\",\n",
        "  # US\n",
        "  \"en-US-Standard-B\", \"en-US-Standard-C\", \"en-US-Standard-D\", \"en-US-Standard-E\",\n",
        "]\n",
        "\n",
        "wavenet_voices = [\n",
        "  # Australia\n",
        "  \"en-AU-Wavenet-A\", \"en-AU-Wavenet-B\", \"en-AU-Wavenet-C\", \"en-AU-Wavenet-D\",\n",
        "  # UK\n",
        "  \"en-GB-Wavenet-A\", \"en-GB-Wavenet-B\", \"en-GB-Wavenet-C\", \"en-GB-Wavenet-D\",\n",
        "  # India\n",
        "  \"en-IN-Wavenet-A\", \"en-IN-Wavenet-B\", \"en-IN-Wavenet-C\",\n",
        "  # US\n",
        "  \"en-US-Wavenet-A\", \"en-US-Wavenet-B\", \"en-US-Wavenet-C\", \"en-US-Wavenet-D\", \"en-US-Wavenet-E\", \"en-US-Wavenet-F\",\n",
        "]\n",
        "\n",
        "cloud_voices = standard_voices + wavenet_voices\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "jey6X7Ga-qvC"
      },
      "outputs": [],
      "source": [
        "#@title Setup Text-to-Speech synthesis\n",
        "\n",
        "from google.cloud import texttospeech\n",
        "\n",
        "# Instantiates a client\n",
        "client = texttospeech.TextToSpeechClient()\n",
        "\n",
        "SAMPLE_RATE = 16000\n",
        "def synthesize_speech(client, text, cloud_voice, language='en-US'):\n",
        "  # Select the type of audio file you want returned\n",
        "  audio_config = texttospeech.types.AudioConfig(\n",
        "      audio_encoding=texttospeech.enums.AudioEncoding.LINEAR16,\n",
        "      sample_rate_hertz=SAMPLE_RATE)\n",
        "\n",
        "  # Set the text input to be synthesized\n",
        "  synthesis_input = texttospeech.types.SynthesisInput(text=text)\n",
        "  voice = texttospeech.types.VoiceSelectionParams(\n",
        "      language_code=language,\n",
        "      name=cloud_voice)\n",
        "\n",
        "\n",
        "  # Perform the text-to-speech request on the text input with the selected\n",
        "  # voice parameters and audio file type\n",
        "  response = client.synthesize_speech(synthesis_input, voice, audio_config)\n",
        "\n",
        "  return response\n",
        "\n",
        "\n",
        "\n",
        "from google.colab import files\n",
        "from io import BytesIO\n",
        "import scipy.io.wavfile as sciwav\n",
        "import collections\n",
        "import time\n",
        "\n",
        "def synthesize_texts(voices, texts, target_words):\n",
        "  cloud_wavs = collections.defaultdict(list)\n",
        "  progress_bar = display(progress(0, 100), display_id=True)\n",
        "  num_synthesized_examples = 0\n",
        "  total_exmaples_to_synthesize = len(voices) * len(texts)\n",
        "  for voice in voices:\n",
        "    for text in texts:\n",
        "      response = synthesize_speech(client, text, voice)\n",
        "      id = '{}_{}'.format(text, voice)\n",
        "      filename = '{}.wav'.format(id)\n",
        "      with open(filename, 'wb') as f_out:\n",
        "        f_out.write(response.audio_content)\n",
        "        # print(\"Audio content written to file '{}'.\".format(filename))\n",
        "        text = text.lower()\n",
        "        text = text if text in target_words else 'negative'\n",
        "        cloud_wavs[text].append(filename)\n",
        "      time.sleep(0.200)\n",
        "      if progress is not None:\n",
        "        num_synthesized_examples += 1\n",
        "        if num_synthesized_examples % 5 == 0:\n",
        "          progress_bar.update(\n",
        "              progress(100 * num_synthesized_examples/total_exmaples_to_synthesize))\n",
        "  progress_bar.update(\n",
        "      progress(100 * num_synthesized_examples/total_exmaples_to_synthesize))\n",
        "  return cloud_wavs\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "krl6d7RI-lKG"
      },
      "outputs": [],
      "source": [
        "#@title Synthesize speech examples\n",
        "\n",
        "%%time\n",
        "\n",
        "print(\"Synthesizing speech examples, this may take a few minutes.\")\n",
        "\n",
        "cloud_wavs = synthesize_texts(\n",
        "    cloud_voices, ALL_WORDS.split(','), TARGET_WORDS.split(','))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gR6n9PMGIHSv"
      },
      "source": [
        "## Load the embedding model\n",
        "\n",
        "The following info messages can be ignored\n",
        "\n",
        "\u003e *INFO:tensorflow:Saver not created because there are no variables in the graph to restore*\n",
        "\n",
        "Don't worry tf hub is restoring all the variables.\n",
        "\n",
        "You can test the model by having it produce an embedding on zeros:\n",
        "\n",
        "\n",
        "```\n",
        "speech_embedding_model.create_embedding(np.zeros((1,66000)))\n",
        "```\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "3TTWa8GN0KHo"
      },
      "outputs": [],
      "source": [
        "#@title Helper functions and classes\n",
        "def normalized_read(filename):\n",
        "  \"\"\"Reads and normalizes a wavfile.\"\"\"\n",
        "  _, data = scipy.io.wavfile.read(open(filename, mode='rb'))\n",
        "  samples_99_percentile = np.percentile(np.abs(data), 99.9)\n",
        "  normalized_samples = data / samples_99_percentile\n",
        "  normalized_samples = np.clip(normalized_samples, -1, 1)\n",
        "  return normalized_samples\n",
        "\n",
        "class EmbeddingDataFileList(object):\n",
        "  \"\"\"Container that loads audio, stores it as embeddings and can\n",
        "  rebalance it.\"\"\"\n",
        "\n",
        "  def __init__(self, filelist,\n",
        "               data_dest_dir,\n",
        "               targets=None,\n",
        "               label_max=10000,\n",
        "               negative_label=\"negative\",\n",
        "               negative_multiplier=25,\n",
        "               target_samples=32000,\n",
        "               progress_bar=None,\n",
        "               embedding_model=None):\n",
        "    \"\"\"Creates an instance of `EmbeddingDataFileList`.\"\"\"\n",
        "    self._negative_label = negative_label\n",
        "    self._data_per_label = collections.defaultdict(list)\n",
        "    self._labelcounts = {}\n",
        "    self._label_list = targets\n",
        "    total_examples = sum([min(len(x), label_max) for x in filelist.values()])\n",
        "    total_examples -= min(len(filelist[negative_label]), label_max)\n",
        "    total_examples += min(len(filelist[negative_label]), negative_multiplier * label_max)\n",
        "    print(\"loading %d examples\" % total_examples)\n",
        "    example_count = 0\n",
        "    for label in filelist:\n",
        "      if label not in self._label_list:\n",
        "        raise ValueError(\"Unknown label:\", label)\n",
        "      label_files = filelist[label]\n",
        "      random.shuffle(label_files)\n",
        "      if label == negative_label:\n",
        "        multplier = negative_multiplier\n",
        "      else:\n",
        "        multplier = 1\n",
        "      for wav_file in label_files[:label_max * multplier]:\n",
        "        data = normalized_read(os.path.join(data_dest_dir, wav_file))\n",
        "        required_padding = target_samples - data.shape[0]\n",
        "        if required_padding \u003e 0:\n",
        "          data = np.pad(data, (required_padding, required_padding), 'constant')\n",
        "        self._labelcounts[label] = self._labelcounts.get(label, 0) + 1\n",
        "        if embedding_model:\n",
        "           data = embedding_model.create_embedding(data)[0][0,:,:,:]\n",
        "        self._data_per_label[label].append(data)\n",
        "        if progress_bar is not None:\n",
        "          example_count += 1\n",
        "          progress_bar.update(progress(100 * example_count/total_examples))\n",
        "\n",
        "  @property\n",
        "  def labels(self):\n",
        "    return self._label_list\n",
        "\n",
        "  def get_label(self, idx):\n",
        "    return self.labels.index(idx)\n",
        "\n",
        "  def _get_filtered_data(self, label, filter_fn):\n",
        "    idx = self.labels.index(label)\n",
        "    return [(filter_fn(x), idx) for x in self._data_per_label[label]]\n",
        "\n",
        "  def _multply_data(self, data, factor):\n",
        "    samples = int((factor - math.floor(factor)) * len(data))\n",
        "    return int(factor) * data + random.sample(data, samples)\n",
        "\n",
        "  def full_rebalance(self, negatives, labeled):\n",
        "    \"\"\"Rebalances for a given ratio of labeled to negatives.\"\"\"\n",
        "    negative_count = self._labelcounts[self._negative_label]\n",
        "    labeled_count = sum(self._labelcounts[key]\n",
        "                        for key in self._labelcounts.keys()\n",
        "                        if key != self._negative_label)\n",
        "    labeled_multiply = labeled * negative_count / (negatives * labeled_count)\n",
        "    for label in self._data_per_label:\n",
        "      if label == self._negative_label:\n",
        "        continue\n",
        "      self._data_per_label[label] = self._multply_data(\n",
        "          self._data_per_label[label], labeled_multiply)\n",
        "      self._labelcounts[label] = len(self._data_per_label[label])\n",
        "\n",
        "  def get_all_data_shuffled(self, filter_fn):\n",
        "    \"\"\"Returns a shuffled list containing all the data.\"\"\"\n",
        "    return self.get_all_data(filter_fn, shuffled=True)\n",
        "\n",
        "  def get_all_data(self, filter_fn, shuffled=False):\n",
        "    \"\"\"Returns a list containing all the data.\"\"\"\n",
        "    data = []\n",
        "    for label in self._data_per_label:\n",
        "      data += self._get_filtered_data(label, filter_fn)\n",
        "    if shuffled:\n",
        "      random.shuffle(data)\n",
        "    return data\n",
        "\n",
        "def cut_middle_frame(embedding, num_frames, flatten):\n",
        "  \"\"\"Extrats the middle frames for an embedding.\"\"\"\n",
        "  left_context = (embedding.shape[0] - num_frames) // 2\n",
        "  if flatten:\n",
        "    return embedding[left_context:left_context+num_frames].flatten()\n",
        "  else:\n",
        "    return embedding[left_context:left_context+num_frames]\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "QLccWY0bIJs1"
      },
      "outputs": [],
      "source": [
        "#@title TfHubWrapper Class\n",
        "\n",
        "class TfHubWrapper(object):\n",
        "  \"\"\"A loads a tf hub embedding model.\"\"\"\n",
        "  def __init__(self, embedding_model_dir):\n",
        "    \"\"\"Creates a `SavedModelWraper`.\"\"\"\n",
        "    self._graph = tf.Graph()\n",
        "    self._sess = tf.Session(graph=self._graph)\n",
        "    with self._graph.as_default():\n",
        "      with self._sess.as_default():\n",
        "        module_spec = hub.load_module_spec(embedding_model_dir)\n",
        "        embedding_module = hub.Module(module_spec)\n",
        "        self._samples = tf.placeholder(\n",
        "            tf.float32, shape=[1, None], name='audio_samples')\n",
        "        self._embedding = embedding_module(self._samples)\n",
        "        self._sess.run(tf.global_variables_initializer())\n",
        "    print(\"Embedding model loaded, embedding shape:\", self._embedding.shape)\n",
        "\n",
        "  def create_embedding(self, samples):\n",
        "    samples = samples.reshape((1, -1))\n",
        "    output = self._sess.run(\n",
        "        [self._embedding],\n",
        "        feed_dict={self._samples: samples})\n",
        "    return output\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "O4LRwMdsIVpo"
      },
      "outputs": [],
      "source": [
        "embedding_model_url = \"https://tfhub.dev/google/speech_embedding/1\"\n",
        "speech_embedding_model = TfHubWrapper(embedding_model_url)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "7jM-7SujSp3T"
      },
      "source": [
        "## Load training and eval data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "hx3Lq7UzP_92"
      },
      "outputs": [],
      "source": [
        "#@title Load evaluation data\n",
        "progress_bar = display(progress(0, 100), display_id=True)\n",
        "\n",
        "print(\"loading eval data\")\n",
        "eval_data = EmbeddingDataFileList(\n",
        "    all_eval_example_files, data_dest_dir,\n",
        "    targets=MODEL_LABELS, embedding_model=speech_embedding_model,\n",
        "    progress_bar=progress_bar)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "eKqMINtnQIf2"
      },
      "outputs": [],
      "source": [
        "#@title Load synthesized examples wav files for training\n",
        "\n",
        "progress_bar = display(progress(0, 100), display_id=True)\n",
        "\n",
        "print(\"loading train data\")\n",
        "train_data = EmbeddingDataFileList(\n",
        "    cloud_wavs, '',\n",
        "    targets=MODEL_LABELS,\n",
        "    embedding_model=speech_embedding_model,\n",
        "    progress_bar=progress_bar)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ipFhC0H14Q9v"
      },
      "source": [
        "## Train and Evaluate a Head Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "uGjxofKb07bk"
      },
      "outputs": [],
      "source": [
        "#@title HeadTrainerClass and head model functions\n",
        "\n",
        "def _fully_connected_model_fn(embeddings, num_labels):\n",
        "  \"\"\"Builds the head model and adds a fully connected output layer.\"\"\"\n",
        "  net = tf.layers.flatten(embeddings)\n",
        "  logits = tf.compat.v1.layers.dense(net, num_labels, activation=None)\n",
        "  return logits\n",
        "\n",
        "framework = tf.contrib.framework\n",
        "layers = tf.contrib.layers\n",
        "\n",
        "def _conv_head_model_fn(embeddings, num_labels, context):\n",
        "  \"\"\"Builds the head model and adds a fully connected output layer.\"\"\"\n",
        "  activation_fn = tf.nn.elu\n",
        "  normalizer_fn = functools.partial(\n",
        "      layers.batch_norm, scale=True, is_training=True)\n",
        "  with framework.arg_scope([layers.conv2d], biases_initializer=None,\n",
        "                           activation_fn=None, stride=1, padding=\"SAME\"):\n",
        "    net = embeddings\n",
        "    net = layers.conv2d(net, 96, [3, 1])\n",
        "    net = normalizer_fn(net)\n",
        "    net = activation_fn(net)\n",
        "    net = layers.max_pool2d(net, [2, 1], stride=[2, 1], padding=\"VALID\")\n",
        "    context //= 2\n",
        "    net = layers.conv2d(net, 96, [3, 1])\n",
        "    net = normalizer_fn(net)\n",
        "    net = activation_fn(net)\n",
        "    net = layers.max_pool2d(net, [context, net.shape[2]], padding=\"VALID\")\n",
        "  net = tf.layers.flatten(net)\n",
        "  logits = layers.fully_connected(\n",
        "      net, num_labels, activation_fn=None)\n",
        "  return logits\n",
        "\n",
        "class HeadTrainer(object):\n",
        "  \"\"\"A tensorflow classifier to quickly train and test on embeddings.\n",
        "\n",
        "  Only use this if you are training a very small model on a very limited amount\n",
        "  of data. If you expect the training to take any more than 15 - 20 min then use\n",
        "  something else.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, model_fn, input_shape, num_targets,\n",
        "               head_learning_rate=0.001, batch_size=64):\n",
        "    \"\"\"Creates a `HeadTrainer`.\n",
        "\n",
        "    Args:\n",
        "      model_fn: function that builds the tensorflow model, defines its loss\n",
        "          and returns the tuple (predictions, loss, accuracy).\n",
        "      input_shape: describes the shape of the models input feature.\n",
        "          Does not include a the batch dimension.\n",
        "      num_targets: Target number of keywords.\n",
        "    \"\"\"\n",
        "    self._input_shape = input_shape\n",
        "    self._output_dim = num_targets\n",
        "    self._batch_size = batch_size\n",
        "    self._graph = tf.Graph()\n",
        "    with self._graph.as_default():\n",
        "      self._feature = tf.placeholder(tf.float32, shape=([None] + input_shape))\n",
        "      self._labels = tf.placeholder(tf.int64, shape=(None))\n",
        "      module_spec = hub.create_module_spec(\n",
        "          module_fn=self._get_headmodule_fn(model_fn, num_targets))\n",
        "      self._module = hub.Module(module_spec, trainable=True)\n",
        "      logits = self._module(self._feature)\n",
        "      self._predictions = tf.nn.softmax(logits)\n",
        "      self._loss, self._accuracy = self._get_loss(\n",
        "          logits, self._labels, self._predictions)\n",
        "      self._update_weights = tf.train.AdamOptimizer(\n",
        "          learning_rate=head_learning_rate).minimize(self._loss)\n",
        "    self._sess = tf.Session(graph=self._graph)\n",
        "    with self._sess.as_default():\n",
        "      with self._graph.as_default():\n",
        "        self._sess.run(tf.local_variables_initializer())\n",
        "        self._sess.run(tf.global_variables_initializer())\n",
        "\n",
        "  def _get_headmodule_fn(self, model_fn, num_targets):\n",
        "    \"\"\"Wraps the model_fn in a tf hub module.\"\"\"\n",
        "    def module_fn():\n",
        "      embeddings = tf.placeholder(\n",
        "          tf.float32, shape=([None] + self._input_shape))\n",
        "      logit = model_fn(embeddings, num_targets)\n",
        "      hub.add_signature(name='default', inputs=embeddings, outputs=logit)\n",
        "    return module_fn\n",
        "\n",
        "\n",
        "  def _get_loss(self, logits, labels, predictions):\n",
        "    \"\"\"Defines the model's loss and accuracy.\"\"\"\n",
        "    xentropy_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
        "        logits=logits, labels=labels)\n",
        "    loss = tf.reduce_mean(xentropy_loss)\n",
        "    accuracy = tf.contrib.metrics.accuracy(tf.argmax(predictions, 1), labels)\n",
        "    return loss, accuracy\n",
        "\n",
        "  def save_head_model(self, save_directory):\n",
        "    \"\"\"Saves the model.\"\"\"\n",
        "    with self._graph.as_default():\n",
        "      self._module.export(save_directory, self._sess)\n",
        "\n",
        "\n",
        "  def _feature_transform(self, batch_features, batch_labels):\n",
        "    \"\"\"Transforms lists of features and labels into into model inputs.\"\"\"\n",
        "    return np.stack(batch_features), np.stack(batch_labels)\n",
        "\n",
        "  def _batch_data(self, data, batch_size=None):\n",
        "    \"\"\"Splits the input data into batches.\"\"\"\n",
        "    batch_features = []\n",
        "    batch_labels = []\n",
        "    batch_size = batch_size or len(data)\n",
        "    for feature, label in data:\n",
        "      if feature.shape != tuple(self._input_shape):\n",
        "        raise ValueError(\n",
        "            \"Feature shape ({}) doesn't match model shape ({})\".format(\n",
        "                feature.shape, self._input_shape))\n",
        "      if not 0 \u003c= label \u003c self._output_dim:\n",
        "\n",
        "        raise ValueError('Label value ({}) outside of target range'.format(\n",
        "            label))\n",
        "      batch_features.append(feature)\n",
        "      batch_labels.append(label)\n",
        "      if len(batch_features) == batch_size:\n",
        "        yield self._feature_transform(batch_features, batch_labels)\n",
        "        del batch_features[:]\n",
        "        del batch_labels[:]\n",
        "    if batch_features:\n",
        "      yield self._feature_transform(batch_features, batch_labels)\n",
        "\n",
        "  def epoch_train(self, data, epochs=1, batch_size=None):\n",
        "    \"\"\"Trains the model on the provided data.\n",
        "\n",
        "    Args:\n",
        "      data: List of tuples (feature, label) where feature is a np array of\n",
        "          shape `self._input_shape` and label an int less than self._output_dim.\n",
        "      epochs: Number of times this data should be trained on.\n",
        "      batch_size: Number of feature, label pairs per batch. Overwrites\n",
        "          `self._batch_size` when set.\n",
        "\n",
        "    Returns:\n",
        "      tuple of accuracy, loss;\n",
        "          accuracy: Average training accuracy.\n",
        "          loss: Loss of the final batch.\n",
        "    \"\"\"\n",
        "    batch_size = batch_size or self._batch_size\n",
        "    accuracy_list = []\n",
        "    for _ in range(epochs):\n",
        "      for features, labels in self._batch_data(data, batch_size):\n",
        "        loss, accuracy, _ = self._sess.run(\n",
        "            [self._loss, self._accuracy, self._update_weights],\n",
        "            feed_dict={self._feature: features, self._labels: labels})\n",
        "        accuracy_list.append(accuracy)\n",
        "    return (sum(accuracy_list))/len(accuracy_list), loss\n",
        "\n",
        "  def test(self, data, batch_size=None):\n",
        "    \"\"\"Evaluates the model on the provided data.\n",
        "\n",
        "    Args:\n",
        "      data: List of tuples (feature, label) where feature is a np array of\n",
        "          shape `self._input_shape` and label an int less than self._output_dim.\n",
        "      batch_size: Number of feature, label pairs per batch. Overwrites\n",
        "          `self._batch_size` when set.\n",
        "\n",
        "    Returns:\n",
        "      tuple of accuracy, loss;\n",
        "          accuracy: Average training accuracy.\n",
        "          loss: Loss of the final batch.\n",
        "    \"\"\"\n",
        "    batch_size = batch_size or self._batch_size\n",
        "    accuracy_list = []\n",
        "    for features, labels in self._batch_data(data, batch_size):\n",
        "      loss, accuracy = self._sess.run(\n",
        "          [self._loss, self._accuracy],\n",
        "          feed_dict={self._feature: features, self._labels: labels})\n",
        "      accuracy_list.append(accuracy)\n",
        "    return sum(accuracy_list)/len(accuracy_list), loss\n",
        "\n",
        "  def infer(self, example_feature):\n",
        "    \"\"\"Runs inference on example_feature.\"\"\"\n",
        "    if example_feature.shape != tuple(self._input_shape):\n",
        "      raise ValueError(\n",
        "          \"Feature shape ({}) doesn't match model shape ({})\".format(\n",
        "              example_feature.shape, self._input_shape))\n",
        "    return self._sess.run(\n",
        "        self._predictions,\n",
        "        feed_dict={self._feature: np.expand_dims(example_feature, axis=0)})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "iUhwy_1bCg3P"
      },
      "outputs": [],
      "source": [
        "#@title Rebalance and filter data.\n",
        "\n",
        "labeled_weight = 8 #@param {type:\"slider\", min:1, max:25, step:1}\n",
        "negatives_weight = 1 #@param {type:\"slider\", min:1, max:25, step:1}\n",
        "context_size = 16 #@param {type:\"slider\", min:1, max:28, step:1}\n",
        "\n",
        "\n",
        "filter_fn = functools.partial(cut_middle_frame, num_frames=context_size, flatten=False)\n",
        "\n",
        "\n",
        "eval_data.full_rebalance(negatives=negatives_weight, labeled=labeled_weight)\n",
        "all_eval_data = eval_data.get_all_data_shuffled(filter_fn=filter_fn)\n",
        "\n",
        "train_data.full_rebalance(negatives=negatives_weight, labeled=labeled_weight)\n",
        "all_train_data = train_data.get_all_data_shuffled(filter_fn=filter_fn)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "bqtJk_XN9Tcb"
      },
      "outputs": [],
      "source": [
        "#@title Run training and evaluation\n",
        "head_model = \"Convolutional\" #@param [\"Convolutional\", \"Fully_Connected\"] {type:\"string\"}\n",
        "\n",
        "#@markdown Suggested **learning_rate** range 0.00001 - 0.01.\n",
        "learning_rate = 0.002 #@param {type:\"number\"}\n",
        "batch_size = 64\n",
        "#@markdown **epochs_per_eval** and **train_eval_loops** control how long the\n",
        "#@markdown the model is trained. An epoch is defined as the model having seen\n",
        "#@markdown each example at least once, with some examples twice to ensure the\n",
        "#@markdown correct labeled / negatives balance.\n",
        "\n",
        "epochs_per_eval = 1 #@param {type:\"slider\", min:1, max:15, step:1}\n",
        "train_eval_loops = 15 #@param {type:\"slider\", min:5, max:80, step:5}\n",
        "\n",
        "if head_model == \"Convolutional\":\n",
        "  model_fn = functools.partial(_conv_head_model_fn, context=context_size)\n",
        "else:\n",
        "  model_fn = _fully_connected_model_fn\n",
        "\n",
        "\n",
        "trainer = HeadTrainer(model_fn=model_fn,\n",
        "                      input_shape=[context_size,1,96],\n",
        "                      num_targets=len(MODEL_LABELS),\n",
        "                      head_learning_rate=learning_rate,\n",
        "                      batch_size=batch_size)\n",
        "\n",
        "data_trained_on = 0\n",
        "data = [] \n",
        "train_results = []\n",
        "eval_results = []\n",
        "max_data = len(all_train_data) * epochs_per_eval * train_eval_loops + 10\n",
        "\n",
        "\n",
        "def plot_step(plot, max_data, data, train_results, eval_results):\n",
        "  plot.clf()\n",
        "  plot.xlim(0, max_data)\n",
        "  plot.ylim(0.85, 1.05)\n",
        "  plot.plot(data, train_results, \"bo\")\n",
        "  plot.plot(data, train_results, \"b\", label=\"train_results\")\n",
        "  if eval_results:\n",
        "    plot.plot(data, eval_results, \"ro\")\n",
        "    plot.plot(data, eval_results, \"r\", label=\"eval_results\")\n",
        "  plot.legend(loc='lower right', fontsize=24)\n",
        "  plot.xlabel('number of examples trained on', fontsize=22)\n",
        "  plot.ylabel('Accuracy', fontsize=22)\n",
        "  plot.xticks(fontsize=20)\n",
        "  plot.yticks(fontsize=20)\n",
        "\n",
        "plt.figure(figsize=(25, 7))\n",
        "for loop in range(train_eval_loops):\n",
        "  train_accuracy, loss = trainer.epoch_train(all_train_data,\n",
        "                                             epochs=epochs_per_eval)\n",
        "  train_results.append(train_accuracy)\n",
        "  if all_eval_data:\n",
        "    eval_accuracy, loss = trainer.test(all_eval_data)\n",
        "    eval_results.append(eval_accuracy)\n",
        "  else:\n",
        "    eval_results = None\n",
        "\n",
        "  data_trained_on += len(all_train_data) * epochs_per_eval\n",
        "  data.append(data_trained_on)\n",
        "  plot_step(plt, max_data, data, train_results, eval_results)\n",
        "\n",
        "  IPython.display.display(plt.gcf())\n",
        "  if all_eval_data:\n",
        "    print(\"Highest eval accuracy: %.2f percent.\" % (100 * max(eval_results)))\n",
        "  IPython.display.clear_output(wait=True)\n",
        "\n",
        "if all_eval_data:\n",
        "  print(\"Highest eval accuracy: %.2f percent.\" % (100 * max(eval_results)))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "LbVy_h8LnlyN"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "Training KWS model with synthesized speech.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
